#include <fstream>
#include <random>
#include <thread>

#include <bthread/bthread.h>
#include <glog/logging.h>
#include <gtest/gtest.h>

#include "hawking/util/node_content.pb.h"
#include "doc/document_creator.h"
#include "index/attribute_index/attribute_disk_indexer/single_value_attribute_disk_indexer.h"
#include "index/attribute_index/attribute_disk_indexer/multi_value_attribute_disk_indexer.h"
#include "index/attribute_index/attribute_mem_indexer/single_value_attribute_mem_indexer.h"
#include "index/attribute_index/attribute_mem_indexer/multi_value_attribute_mem_indexer.h"
#include "util/file_wrapper/file_wrapper.h"

namespace hawking {
namespace indexlib {

DECLARE_int32(io_uring_queue_size);

}
}

using namespace hawking::indexlib;

std::uniform_int_distribution<int> dis(1, 1024000);
std::uniform_int_distribution<int> dis2(1, 33);
std::default_random_engine e{std::random_device{}()};

TEST(load_test, single) {
    SingleValueAttributeDiskIndexer<int64_t> disk_indexer;
    Status ret = disk_indexer.Load(
        "/tmp/single_dump",
        hawking::util::ColumnType::COLUMN_INT64,
        FSFileType::FSFT_BLOCK);
    EXPECT_EQ(ret, Status::OK);
}

TEST(func_test, single_function) {
    SingleValueAttributeMemIndexer<int64_t> single_attr_mem_indexer;
    single_attr_mem_indexer.Init(hawking::util::COLUMN_INT64, 0);

    // write mem indexer
    std::unique_ptr<NormalDocument> doc = std::make_unique<NormalDocument>();
    std::shared_ptr<AttributeDocument> attr_doc = std::make_shared<AttributeDocument>();
    attr_doc->Reserve(1);
    doc->SetAttributeDocument(attr_doc);
    doc->SetOperatorType(hawking::util::IndexOp::INDEX_RT_ADD);
    SingleValueAttributeConvertor<int64_t> convertor(hawking::util::COLUMN_INT64);
    std::unordered_map<DocId, DocId> old2new_map;
    for (DocId idx = 0; idx < 1024000; ++idx) {
        hawking::util::FieldValue value;
        value.set_int_value(idx);
        std::string_view sv = convertor.Encode(&value, single_attr_mem_indexer.GetPool());

        attr_doc->SetField(0, sv);
        doc->SetDocId(idx);
        single_attr_mem_indexer.AddDocument(doc.get());
        old2new_map.emplace(idx, idx);
    }

    // mem read
    for (DocId idx = 0; idx < 1024000; ++idx) {
        hawking::util::FieldValue field_value;
        single_attr_mem_indexer.Read(idx, &field_value);
        EXPECT_EQ(field_value.int_value(), idx);
    }

    // dump
    FslibWrapper::Mkdir("/tmp/single_dump");
    single_attr_mem_indexer.Dump("/tmp/single_dump", old2new_map);

    // load
    SingleValueAttributeDiskIndexer<int64_t> disk_indexer;
    Status ret = disk_indexer.Load(
        "/tmp/single_dump",
        hawking::util::ColumnType::COLUMN_INT64,
        FSFileType::FSFT_BLOCK);
    EXPECT_EQ(ret, Status::OK);

    // disk read
    int64_t ret_value;
    for (DocId idx = 0; idx < 1024000; ++idx) {
        disk_indexer.Read(
            (uint8_t*)&ret_value,
            disk_indexer.ElementLength(),
            disk_indexer.GetOffset(idx));
        EXPECT_EQ(ret_value, idx);
    }

    // disk batch read
    std::vector<int64_t> ret_values;
    std::vector<DocId> verify_values;
    verify_values.reserve(100);
    ret_values.resize(100);
    std::vector<FilePart> file_parts;
    file_parts.reserve(100);
    for (size_t idx = 0; idx < 100; ++idx) {
        DocId random_doc_id = dis(e);
        verify_values.push_back(random_doc_id);
        file_parts.emplace_back(
            disk_indexer.GetOffset(random_doc_id),
            disk_indexer.ElementLength(),
            (uint8_t*)&ret_values[idx]);
    }

    disk_indexer.BatchRead(file_parts);
    for (size_t idx = 0; idx < 100; ++idx) {
        EXPECT_EQ(ret_values[idx], verify_values[idx]);
    }
}

TEST(load_test, string) {
    SingleValueAttributeDiskIndexer<int64_t> disk_indexer;
    Status ret = disk_indexer.Load(
        "/tmp/string_dump",
        hawking::util::ColumnType::COLUMN_INT64,
        FSFileType::FSFT_BLOCK);
    EXPECT_EQ(ret, Status::OK);
}

TEST(func_test, string_function) {
#if 0
    MultiValueAttributeMemIndexer<char> string_attr_mem_indexer;
    string_attr_mem_indexer.Init(hawking::util::COLUMN_STRING, 0);

    // write mem indexer
    std::unique_ptr<NormalDocument> doc = std::make_unique<NormalDocument>();
    std::shared_ptr<AttributeDocument> attr_doc = std::make_shared<AttributeDocument>();
    attr_doc->Reserve(1);
    doc->SetAttributeDocument(attr_doc);
    doc->SetOperatorType(hawking::util::IndexOp::INDEX_RT_ADD);
    StringAttributeConvertor convertor(hawking::util::COLUMN_STRING);
    std::unordered_map<DocId, DocId> old2new_map;
    for (DocId idx = 0; idx < 51200000; ++idx) {
        hawking::util::FieldValue value;
        value.set_bytes_value(std::to_string(idx));
        std::string_view sv = convertor.Encode(&value, string_attr_mem_indexer.GetPool());

        attr_doc->SetField(0, sv);
        doc->SetDocId(idx);
        string_attr_mem_indexer.AddDocument(doc.get());
        old2new_map.emplace(idx, idx);
    }

    // mem read
    for (DocId idx = 0; idx < 51200000; ++idx) {
        hawking::util::FieldValue field_value;
        string_attr_mem_indexer.Read(idx, &field_value);
        EXPECT_EQ(field_value.bytes_value(), std::to_string(idx));
    }

    // dump
    FslibWrapper::Mkdir("/tmp/string_dump");
    string_attr_mem_indexer.Dump("/tmp/string_dump", old2new_map);
#endif

    // load
    MultiValueAttributeDiskIndexer<char> disk_indexer;
    Status ret = disk_indexer.Load(
        "/tmp/string_dump",
        hawking::util::ColumnType::COLUMN_STRING,
        FSFileType::FSFT_BLOCK);
    EXPECT_EQ(ret, Status::OK);

    // disk read
    hawking::util::ArenaWrapper arena_wrapper;
    SegmentContext segment_ctx(arena_wrapper);
    std::vector<std::string> verify_ret_values;
    verify_ret_values.reserve(2000);
    segment_ctx.Reserve(2000);
    for (DocId idx = 0; idx < 2000; ++idx) {
        DocId random_doc_id = dis(e);
        verify_ret_values.push_back(std::to_string(random_doc_id));
        segment_ctx.PushDocId(random_doc_id);
    }
    std::vector<std::pair<uint64_t, ArrayLengthType>> array_info;
    array_info.reserve(2000);
    disk_indexer.GetCountAndElemLens(&segment_ctx, &array_info);
    // for (const auto& info : array_info) {
    //     LOG(INFO) << "offset " << info.first << ", string length "
    //         << info.second;
    // }

    for (DocId idx = 0; idx < 2000; ++idx) {
        auto& info = array_info[idx];
        std::string ret_value;
        ret_value.resize(info.second);
        disk_indexer.Read((uint8_t*)ret_value.data(), info.second, info.first);
        EXPECT_EQ(verify_ret_values[idx], ret_value);
    }
    
    // disk batch read
    std::vector<std::thread> workers;
    workers.reserve(16);
    for (size_t idx = 0; idx < workers.capacity(); ++idx) {
        workers.emplace_back([&disk_indexer](){
            while (1) {
                std::vector<FilePart> datafile_part;
                hawking::util::ArenaWrapper arena_wrapper;
                SegmentContext segment_ctx(arena_wrapper);
                std::vector<std::pair<uint64_t, ArrayLengthType>> array_info;
                segment_ctx.Reserve(2000);
                for (size_t idx = 0; idx < 2000; ++idx) {
                    DocId random_doc_id = dis(e);
                    segment_ctx.PushDocId(random_doc_id);
                }
            
                std::chrono::high_resolution_clock::time_point beginTime = std::chrono::high_resolution_clock::now();
            
            
                disk_indexer.GetCountAndElemLens(&segment_ctx, &array_info);
            
                std::vector<DocContext>& doc_context = *(segment_ctx.MutableDocContext());
                for (size_t idx = 0; idx < 2000; ++idx) {
                    const auto& info = array_info[idx];
                    DocContext& doc_ctx = doc_context[idx];
            
                    hawking::util::FieldValue* field_value =
                        segment_ctx.ArenaCtx().CreateOnArena<hawking::util::FieldValue>();
                    doc_ctx.field2attr.emplace(0, field_value);
                    field_value->mutable_bytes_value()->resize(info.second);
            
                    datafile_part.emplace_back(
                        info.first,
                        info.second,
                        reinterpret_cast<uint8_t*>(field_value->mutable_bytes_value()->data()));
                }
            
                std::sort(datafile_part.begin(), datafile_part.end());
                disk_indexer.BatchRead(datafile_part);
                std::chrono::high_resolution_clock::time_point endTime = std::chrono::high_resolution_clock::now();
                std::chrono::milliseconds timeInterval = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - beginTime);
                LOG(INFO) << "cost " << timeInterval.count() << "ms";
                for (size_t idx = 0; idx < 2000; ++idx) {
                    DocContext& doc_ctx = doc_context[idx];
                    hawking::util::FieldValue* field_value = doc_ctx.GetField(0);
                    // LOG(INFO) << field_value->bytes_value();
                    EXPECT_EQ(field_value->bytes_value(), std::to_string(segment_ctx.DocIds()[idx]));
                }
             }
        });
    }

    for (auto& worker : workers) {
        worker.join();
    }
}

TEST(load_test, multi_value) {
#if 0
    MultiValueAttributeMemIndexer<float> float_attr_mem_indexer;
    float_attr_mem_indexer.Init(hawking::util::COLUMN_FLOAT_LIST, 0);

    // write mem indexer
    std::unique_ptr<NormalDocument> doc = std::make_unique<NormalDocument>();
    std::shared_ptr<AttributeDocument> attr_doc = std::make_shared<AttributeDocument>();
    attr_doc->Reserve(1);
    doc->SetAttributeDocument(attr_doc);
    doc->SetOperatorType(hawking::util::IndexOp::INDEX_RT_ADD);
    MultiValueAttributeConvertor<float> convertor(hawking::util::COLUMN_FLOAT_LIST);
    std::unordered_map<DocId, DocId> old2new_map;
    for (DocId idx = 0; idx < 5120000; ++idx) {
        hawking::util::FieldValue value;
        int32_t elem_count = dis2(e);
        // LOG(INFO) << "idx " << idx << ", elem_count " << elem_count;
        value.mutable_double_values()->mutable_values()->Reserve(elem_count);
        for (int32_t idx = 0; idx < elem_count; ++idx) {
            value.mutable_double_values()->add_values(static_cast<float>(idx));
        }
        std::string_view sv = convertor.Encode(&value, float_attr_mem_indexer.GetPool());

        attr_doc->SetField(0, sv);
        doc->SetDocId(idx);
        float_attr_mem_indexer.AddDocument(doc.get());
        old2new_map.emplace(idx, idx);
    }

    // mem read
    for (DocId idx = 0; idx < 5120000; ++idx) {
        hawking::util::FieldValue field_value;
        float_attr_mem_indexer.Read(idx, &field_value);
        // LOG(INFO) << "idx " << idx << ", mem read elem count " << field_value.double_values().values_size();
        // for (size_t edx = 0; edx < field_value.double_values().values_size(); ++edx) {
        //     LOG(INFO) << "ele " << field_value.double_values().values(edx);
        // }
    }

    // dump
    FslibWrapper::Mkdir("/tmp/float_list_dump");
    float_attr_mem_indexer.Dump("/tmp/float_list_dump", old2new_map);
#endif

    // load
    MultiValueAttributeDiskIndexer<float> disk_indexer;
    Status ret = disk_indexer.Load(
        "/tmp/float_list_dump",
        hawking::util::ColumnType::COLUMN_STRING,
        FSFileType::FSFT_BLOCK);
    EXPECT_EQ(ret, Status::OK);

    // disk read
    hawking::util::ArenaWrapper arena_wrapper;
    SegmentContext segment_ctx(arena_wrapper);
    segment_ctx.Reserve(16);
    for (DocId idx = 0; idx < 16; ++idx) {
        segment_ctx.PushDocId(idx);
    }
    std::vector<std::pair<uint64_t, ArrayLengthType>> array_info;
    array_info.reserve(16);
    disk_indexer.GetCountAndElemLens(&segment_ctx, &array_info);
    // for (const auto& info : array_info) {
    //     LOG(INFO) << "offset " << info.first << ", string length "
    //         << info.second;
    // }

    for (DocId idx = 0; idx < 16; ++idx) {
        auto& info = array_info[idx];
        std::vector<float> array(info.second, 0);
        disk_indexer.Read(
            (uint8_t*)(array.data()),
            info.second * disk_indexer.ElementLength(),
            info.first);
        // LOG(INFO) << "idx " << idx << ", elem_count " << info.second;
        // for (ArrayLengthType edx = 0; edx < info.second; ++edx) {
        //     LOG(INFO) << array[edx];
        // }
    }
    
    // // disk batch read
    std::vector<std::thread> workers;
    workers.reserve(16);
    for (size_t idx = 0; idx < workers.capacity(); ++idx) {
        workers.emplace_back([&disk_indexer](){
            while (1) {
                std::vector<FilePart> datafile_part;
                hawking::util::ArenaWrapper arena_wrapper;
                SegmentContext segment_ctx(arena_wrapper);
                std::vector<std::pair<uint64_t, ArrayLengthType>> array_info;
                segment_ctx.Reserve(2000);
                for (size_t idx = 0; idx < 2000; ++idx) {
                    DocId random_doc_id = dis(e);
                    segment_ctx.PushDocId(random_doc_id);
                    // segment_ctx.PushDocId(idx);
                }
            
                std::chrono::high_resolution_clock::time_point beginTime = std::chrono::high_resolution_clock::now();
            
            
                disk_indexer.GetCountAndElemLens(&segment_ctx, &array_info);
            
                std::vector<DocContext>& doc_context = *(segment_ctx.MutableDocContext());
                for (size_t idx = 0; idx < 2000; ++idx) {
                    const auto& info = array_info[idx];
                    DocContext& doc_ctx = doc_context[idx];
            
                    hawking::util::FieldValue* field_value =
                        segment_ctx.ArenaCtx().CreateOnArena<hawking::util::FieldValue>();
                    doc_ctx.field2attr.emplace(0, field_value);
                    field_value->mutable_double_values()->mutable_values()->Resize(info.second, 0);

                    for (ArrayLengthType edx = 0; edx < info.second; ++edx) {
                        double* p_double = field_value->mutable_double_values()->mutable_values()->Mutable(edx);
                        datafile_part.emplace_back(
                            info.first + edx * disk_indexer.ElementLength(),
                            disk_indexer.ElementLength(),
                            reinterpret_cast<uint8_t*>(p_double));
                    }
                }
            
                std::sort(datafile_part.begin(), datafile_part.end());
                disk_indexer.BatchRead(datafile_part);
                for (size_t idx = 0; idx < 2000; ++idx) {
                    DocContext& doc_ctx = doc_context[idx];
                    hawking::util::FieldValue* field_value = doc_ctx.GetField(0);

                    for (int32_t edx = 0; edx < field_value->double_values().values_size(); ++edx) {
                        float* f_ptr = reinterpret_cast<float*>(field_value->mutable_double_values()->mutable_values()->Mutable(edx));
                        field_value->mutable_double_values()->mutable_values()->Set(edx, *f_ptr);
                    }

                    // LOG(INFO) << "idx " << idx << ", elem_count " << field_value->double_values().values_size();
                    // for (int32_t edx = 0; edx < field_value->double_values().values_size(); ++edx) {
                    //     LOG(INFO) << "ele " << field_value->double_values().values(edx);
                    // }
                }

                std::chrono::high_resolution_clock::time_point endTime = std::chrono::high_resolution_clock::now();
                std::chrono::milliseconds timeInterval = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - beginTime);
                LOG(INFO) << "cost " << timeInterval.count() << "ms";
            }
        });
    }

    for (auto& worker : workers) {
        worker.join();
    }
}

TEST(load_test, multi_string) {
#if 0
    MultiValueAttributeMemIndexer<std::string> string_attr_mem_indexer;
    string_attr_mem_indexer.Init(hawking::util::COLUMN_STRING_LIST, 0);

    // write mem indexer
    std::unique_ptr<NormalDocument> doc = std::make_unique<NormalDocument>();
    std::shared_ptr<AttributeDocument> attr_doc = std::make_shared<AttributeDocument>();
    attr_doc->Reserve(1);
    doc->SetAttributeDocument(attr_doc);
    doc->SetOperatorType(hawking::util::IndexOp::INDEX_RT_ADD);
    MultiStringAttributeConvertor convertor(hawking::util::COLUMN_STRING_LIST);
    std::unordered_map<DocId, DocId> old2new_map;
    for (DocId idx = 0; idx < 5120000; ++idx) {
        hawking::util::FieldValue value;
        int32_t elem_count = dis2(e);
        // LOG(INFO) << "idx " << idx << ", elem_count " << elem_count;
        value.mutable_bytes_values()->mutable_values()->Reserve(elem_count);
        for (int32_t idx = 0; idx < elem_count; ++idx) {
            value.mutable_bytes_values()->add_values(std::to_string(idx));
        }
        std::string_view sv = convertor.Encode(&value, string_attr_mem_indexer.GetPool());

        attr_doc->SetField(0, sv);
        doc->SetDocId(idx);
        string_attr_mem_indexer.AddDocument(doc.get());
        old2new_map.emplace(idx, idx);
    }

    // mem read
    for (DocId idx = 0; idx < 5120000; ++idx) {
        hawking::util::FieldValue field_value;
        string_attr_mem_indexer.Read(idx, &field_value);
        // LOG(INFO) << "idx " << idx << ", mem read elem count " << field_value.bytes_values().values_size();
        // for (size_t edx = 0; edx < field_value.bytes_values().values_size(); ++edx) {
        //     LOG(INFO) << "ele " << field_value.bytes_values().values(edx);
        // }
    }

    // dump
    FslibWrapper::Mkdir("/tmp/string_list_dump");
    string_attr_mem_indexer.Dump("/tmp/string_list_dump", old2new_map);
#endif

    // load
    MultiValueAttributeDiskIndexer<std::string> disk_indexer;
    Status ret = disk_indexer.Load(
        "/tmp/string_list_dump",
        hawking::util::ColumnType::COLUMN_STRING_LIST,
        FSFileType::FSFT_BLOCK);
    EXPECT_EQ(ret, Status::OK);

    // disk read
    hawking::util::ArenaWrapper arena_wrapper;
    SegmentContext segment_ctx(arena_wrapper);
    segment_ctx.Reserve(2000);
    for (DocId idx = 0; idx < 2000; ++idx) {
        segment_ctx.PushDocId(idx);
    }
    std::vector<const std::pair<uint64_t, std::vector<StringLengthType>>*> array_info;
    array_info.reserve(2000);
    disk_indexer.GetCountAndElemLens(&segment_ctx, &array_info);
    // for (const auto& info : array_info) {
    //     LOG(INFO) << "offset " << info.first << ", string length "
    //         << info.second;
    // }

    for (DocId idx = 0; idx < 2000; ++idx) {
        auto* info = array_info[idx];
        off_t offset = info->first;
        std::vector<std::string> array(info->second.size(), "");
        // LOG(INFO) << "idx " << idx << ", elem_count " << info->second.size();
        for (size_t edx = 0; edx < info->second.size(); ++edx) {
            array[edx].resize(info->second[edx]);
            disk_indexer.Read(
                (uint8_t*)(array[edx].data()),
                info->second[edx],
                offset);

            offset += info->second[edx];
            // LOG(INFO) << array[edx];
        }
    }
    
    // disk batch read
    std::vector<std::thread> workers;
    workers.reserve(16);
    for (size_t idx = 0; idx < workers.capacity(); ++idx) {
        workers.emplace_back([&disk_indexer](){
            while (1) {
                std::vector<FilePart> datafile_part;
                hawking::util::ArenaWrapper arena_wrapper;
                SegmentContext segment_ctx(arena_wrapper);
                std::vector<const std::pair<uint64_t, std::vector<StringLengthType>>*> array_info;
                segment_ctx.Reserve(2000);
                for (size_t idx = 0; idx < 2000; ++idx) {
                    DocId random_doc_id = dis(e);
                    segment_ctx.PushDocId(random_doc_id);
                    // segment_ctx.PushDocId(idx);
                }
            
                std::chrono::high_resolution_clock::time_point beginTime = std::chrono::high_resolution_clock::now();
            
            
                disk_indexer.GetCountAndElemLens(&segment_ctx, &array_info);
            
                std::vector<DocContext>& doc_context = *(segment_ctx.MutableDocContext());
                for (size_t idx = 0; idx < 2000; ++idx) {
                    const auto& info = *array_info[idx];
                    DocContext& doc_ctx = doc_context[idx];
            
                    hawking::util::FieldValue* field_value =
                        segment_ctx.ArenaCtx().CreateOnArena<hawking::util::FieldValue>();
                    doc_ctx.field2attr.emplace(0, field_value);
                    field_value->mutable_bytes_values()->mutable_values()->Reserve(info.second.size());

                    off_t offset = info.first;
                    for (ArrayLengthType edx = 0; edx < info.second.size(); ++edx) {
                        std::string* p_string = field_value->mutable_bytes_values()->add_values();
                        p_string->resize(info.second[edx]);
                        datafile_part.emplace_back(
                            offset,
                            info.second[edx],
                            reinterpret_cast<uint8_t*>(p_string->data()));
                        offset += info.second[edx];
                    }
                }
            
                std::sort(datafile_part.begin(), datafile_part.end());
                disk_indexer.BatchRead(datafile_part);
                // for (size_t idx = 0; idx < 2000; ++idx) {
                //     DocContext& doc_ctx = doc_context[idx];
                //     hawking::util::FieldValue* field_value = doc_ctx.GetField(0);

                //     LOG(INFO) << "idx " << idx << ", elem_count " << field_value->bytes_values().values_size();
                //     for (int32_t edx = 0; edx < field_value->bytes_values().values_size(); ++edx) {
                //         LOG(INFO) << "ele " << field_value->bytes_values().values(edx);
                //     }
                // }

                std::chrono::high_resolution_clock::time_point endTime = std::chrono::high_resolution_clock::now();
                std::chrono::milliseconds timeInterval = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - beginTime);
                LOG_EVERY_N(INFO, 1000) << "cost " << timeInterval.count() << "ms";
            }
        });
    }

    for (auto& worker : workers) {
        worker.join();
    }
}

int main(int argc, char* argv[]) {
    ::testing::InitGoogleTest(&argc, argv);
    ::google::ParseCommandLineFlags(&argc, &argv, true);
    srand(static_cast<unsigned>(time(0)));
    return RUN_ALL_TESTS();
}
